import json
import datetime
import logging
import asyncio
import nls
from aliyunsdkcore.client import AcsClient
from aliyunsdkcore.request import CommonRequest
from aliyunsdkcore.auth.credentials import AccessKeyCredential

logger = logging.getLogger(__name__)

# 全局并发控制：最多 2 路
_concurrency_semaphore = asyncio.BoundedSemaphore(2)


class TingwuRealtimeSTT:
    """通义听悟实时转写服务（并发限 2 路）；可选启用服务端 VAD 端点检测"""

    def __init__(self, app_key, access_key_id, access_key_secret, region="cn-beijing"):
        self.app_key = app_key
        self.ak = access_key_id
        self.sk = access_key_secret
        self.region = region

        self.task_id = None
        self.ws_url = None
        self.nls_client = None
        self.connected = False

        # 回调
        self.on_transcript = None
        self.on_error = None

        # 主事件循环引用（用于跨线程调度回调）
        self.main_loop = None

        self.logger = logger

    async def create_task(self):
        credentials = AccessKeyCredential(self.ak, self.sk)
        client = AcsClient(region_id=self.region, credential=credentials)

        body = {
            "AppKey": self.app_key,
            "Input": {
                "Format": "pcm",
                "SampleRate": 16000,
                "SourceLanguage": "cn",
                "ProgressiveCallbacksEnabled": True,
                "TaskKey": f"task_{datetime.datetime.now():%Y%m%d%H%M%S}"
            },
            "Parameters": {
                "InterimResultEnable": True,
                "PunctuationEnable": True,
                # 关闭服务端静默检测，本地 VAD 控制
                "TurnDetection": {
                    "Enable": False
                }
            }
        }

        req = CommonRequest()
        req.set_accept_format('json')
        req.set_domain(f'tingwu.{self.region}.aliyuncs.com')
        req.set_version('2023-09-30')
        req.set_protocol_type('https')
        req.set_method('PUT')
        req.set_uri_pattern('/openapi/tingwu/v2/tasks')
        req.add_query_param('type', 'realtime')
        req.add_header('Content-Type', 'application/json')
        req.set_content(json.dumps(body).encode('utf-8'))

        self.logger.info("Creating Tingwu realtime task...")
        # 加超时保护
        resp = await asyncio.wait_for(
            asyncio.get_running_loop()
                .run_in_executor(None, client.do_action_with_exception, req),
            timeout=20.0
        )
        data = json.loads(resp.decode() if isinstance(resp, bytes) else resp)

        code = data.get("Code") or data.get("code")
        msg = data.get("Message", data.get("message", ""))
        success = code in ("success", 0) or data.get("Success", False) or msg.lower() == "success"
        if not success:
            raise RuntimeError(f"CreateTask failed: {msg} (Code: {code})")

        d = data.get("Data", {})
        self.task_id = d.get("TaskId")
        self.ws_url = d.get("WebSocketUrl") or d.get("MeetingJoinUrl")
        if not self.task_id or not self.ws_url:
            raise RuntimeError(f"Invalid task data: {data}")

        self.logger.info(f"Task created: {self.task_id}")
        return self.task_id, self.ws_url

    async def connect(self):
        async with _concurrency_semaphore:
            try:
                await self.create_task()
                self.logger.info(f"Connecting to STT at {self.ws_url}")
                self.nls_client = nls.NlsRealtimeMeeting(
                    url=self.ws_url,
                    on_sentence_begin=self._on_sentence_begin,
                    on_sentence_end=self._on_sentence_end,
                    on_result_changed=self._on_result_changed,
                    on_error=self._on_error,
                    on_close=self._on_close,
                    callback_args=[self]
                )
                await asyncio.get_running_loop().run_in_executor(None, self.nls_client.start)
                self.connected = True
                return True
            except Exception as e:
                self.logger.exception("STT connect failed")
                self.connected = False
                if self.on_error:
                    self.on_error(str(e))
                return False

    async def send_audio(self, pcm_data: bytes, throttle: bool = True):
        if not self.connected or not self.nls_client:
            self.logger.warning("STT not connected, drop audio")
            return False
        if not pcm_data:
            return True
        chunk_size = 640
        try:
            for i in range(0, len(pcm_data), chunk_size):
                chunk = pcm_data[i:i+chunk_size]
                await asyncio.get_running_loop().run_in_executor(None, self.nls_client.send_audio, chunk)
                if throttle:
                    await asyncio.sleep(0.02)
            return True
        except Exception as e:
            self.logger.error(f"send_audio failed: {e}")
            self.connected = False
            if self.on_error:
                self.on_error(str(e))
            return False

    async def close(self):
        # 统一关闭逻辑，不涉及并发信号量
        if self.nls_client:
            try:
                self.nls_client.stop()
                self.logger.info("STT WebSocket stopped")
            except Exception as e:
                self.logger.error(f"Error stopping WS: {e}")
            finally:
                self.nls_client = None
                self.connected = False

        result = {}
        if self.task_id:
            # 停止任务并查询信息
            try:
                credentials = AccessKeyCredential(self.ak, self.sk)
                client = AcsClient(region_id=self.region, credential=credentials)

                req = CommonRequest()
                req.set_accept_format('json')
                req.set_domain(f'tingwu.{self.region}.aliyuncs.com')
                req.set_version('2023-09-30')
                req.set_protocol_type('https')
                req.set_method('PUT')
                req.set_uri_pattern('/openapi/tingwu/v2/tasks')
                req.add_query_param('type', 'realtime')
                req.add_query_param('operation', 'stop')
                req.add_header('Content-Type', 'application/json')
                req.set_content(json.dumps({
                    "AppKey": self.app_key,
                    "Input": {"TaskId": self.task_id}
                }).encode('utf-8'))

                await asyncio.wait_for(
                    asyncio.get_running_loop()
                        .run_in_executor(None, client.do_action_with_exception, req),
                    timeout=5.0
                )
            except Exception as e:
                self.logger.error(f"StopTask failed: {e}")

            # 获取任务结束信息
            result = await self._get_task_info()

        return result

    async def _get_task_info(self):
        if not self.task_id:
            return {}
        try:
            credentials = AccessKeyCredential(self.ak, self.sk)
            client = AcsClient(region_id=self.region, credential=credentials)

            req = CommonRequest()
            req.set_accept_format('json')
            req.set_domain(f'tingwu.{self.region}.aliyuncs.com')
            req.set_version('2023-09-30')
            req.set_protocol_type('https')
            req.set_method('GET')
            req.set_uri_pattern(f'/openapi/tingwu/v2/tasks/{self.task_id}')

            resp = await asyncio.wait_for(
                asyncio.get_running_loop()
                    .run_in_executor(None, client.do_action_with_exception, req),
                timeout=5.0
            )
            data = json.loads(resp.decode() if isinstance(resp, bytes) else resp)
            return data.get('Data', {})
        except Exception as e:
            self.logger.error(f"GetTaskInfo error: {e}")
            return {}

    # —— 回调实现 —— #
    def _parse_message(self, msg):
        if isinstance(msg, str):
            try:
                return json.loads(msg)
            except:
                return {"raw": msg}
        if isinstance(msg, dict):
            return msg
        return {"raw": str(msg)}

    def _on_sentence_begin(self, msg, *args):
        self.logger.debug("Sentence begin")

    def _on_sentence_end(self, msg, *args):
        data = self._parse_message(msg)
        text = (data.get("payload") or {}).get("result") or data.get("text") or ""
        text = text.strip()
        if text and self.on_transcript:
            self.on_transcript(text, True)

    def _on_result_changed(self, msg, *args):
        data = self._parse_message(msg)
        text = (data.get("payload") or {}).get("result") or data.get("text") or ""
        text = text.strip()
        if text and self.on_transcript:
            self.on_transcript(text, False)

    def _on_error(self, msg, *args):
        data = self._parse_message(msg)
        header = data.get("header", {}) or {}
        status = header.get("status", 0)

        # 忽略空音频和静默超时
        if status in (40000000, 40000004):
            level = "warning" if status == 40000000 else "info"
            getattr(self.logger, level)(f"Ignored STT status: {status}")
            if status == 40000004 and self.nls_client:
                self.nls_client.stop()
                self.connected = False
            return

        err = data.get("message") or data.get("error") or str(data)
        self.logger.error(f"STT error: {err}")
        self.connected = False
        if self.on_error:
            self.on_error(err)

    def _on_close(self, *args):
        self.logger.info("STT connection closed")
        self.connected = False
